from torchinfo import summary

from configs.data_configs import get_dataset_class
from configs.hparams import get_hparams_class
from models.lwCETModel import lwCET
from myutils.configUtils import get_configs

# 输出模型的结构


dataset_configs, hparams_class1 = get_configs()
hparams = hparams_class1.train_params
# 创建模型实例
model = lwCET(dataset_configs, hparams,add_fea=True)

# 使用torchsummary输出模型结构
# 假设输入尺寸：图像为 [1, 100, 100]，表格特征为4维
summary(
    model,
    input_size=[(128, 1, 268), (128, 52)],  # 保持原样即可
    col_names=["input_size", "output_size", "num_params", "params_percent", "kernel_size", "mult_adds", "trainable"],
    verbose=2,
    device='cpu'
)
